from flask import Flask, render_template, Response, jsonify
import cv2
import mediapipe as mp
import numpy as np
import PoseModule as pm

app = Flask(__name__)
@app.route('/')
def index_1():
    return render_template('index_1.html')
@app.route('/index.html')
def index():
    return render_template('index.html')
def calculate_body_balance(joints_data):
    # 在这里计算身体重心的位置
    # 这里使用简单的平均值方法来计算重心
    x = sum(j[0] for j in joints_data) / len(joints_data)
    y = sum(j[1] for j in joints_data) / len(joints_data)
    return x, y

def calculate_movement_range(joints_data):
    # 在这里计算动作的幅度范围
    # 这里使用简单的欧几里得距离来计算幅度
    min_x = min(j[0] for j in joints_data)
    max_x = max(j[0] for j in joints_data)
    min_y = min(j[1] for j in joints_data)
    max_y = max(j[1] for j in joints_data)
    return ((max_x - min_x) ** 2 + (max_y - min_y) ** 2) ** 0.5

def check_pushup_range(lmList):
    # 计算动作的幅度范围
    movement_range = calculate_movement_range(lmList)

    # 设置您的判断标准
    min_range = 10  # 最小幅度范围
    max_range = 1000  # 最大幅度范围

    # 判断动作是否达标
    if movement_range >= min_range and movement_range <= max_range:
        return True
    else:
        return False

def calculate_center_of_gravity(lmList):
    # Extract landmark coordinates
    coords = []
    for lm in lmList:
        x, y, _ = lm
        coords.append([x, y])

    # Calculate centroid of the landmarks
    centroid = np.mean(coords, axis=0)

    # Calculate distance of markers from the centroid
    distances = []
    for coord in coords:
        dist = np.linalg.norm(coord - centroid)
        distances.append(dist)

    # Calculate the weight of each marker
    weights = [float(dist / sum(distances)) for dist in distances]  # 强制转换为float型

    # Calculate center of gravity
    cog = np.array([0.0, 0.0])
    for i, coord in enumerate(coords):
        coord = np.array(coord, dtype=float)  # 转换为Numpy数组，并确保元素类型为float
        cog += weights[i] * coord

    return cog

def generate_frames():
    info = {"pushup_count": 0}
    cap = cv2.VideoCapture(0)
    detector = pm.poseDetector()
    count = 0
    direction = 0
    form = 0
    progress_position = 0
    feedback = "Fix Form"

    while cap.isOpened():
        ret, img = cap.read()  # 640 x 480
        # Determine dimensions of video - Help with creation of box in Line 43
        # print(width, height)

        img = detector.findPose(img, False)
        lmList = detector.findPosition(img, False)
        # print(lmList)
        if len(lmList) != 0:
            elbow = detector.findAngle(img, 11, 13, 15)
            shoulder = detector.findAngle(img, 13, 11, 23)
            hip = detector.findAngle(img, 11, 23, 25)

            # Percentage of success of pushup
            per = np.interp(elbow, (90, 160), (0, 100))

            # Bar to show Pushup progress
            bar = np.interp(elbow, (90, 160), (380, 50))

            # Check to ensure right form before starting the program
            if elbow > 160 and shoulder > 40 and hip > 160 and check_pushup_range(lmList):
                form = 1

            # Check for full range of motion for the pushup
            if form == 1:
                if per == 0:
                    if elbow <= 90 and hip > 160:
                        feedback = "Up"
                        if direction == 0:
                            count += 0.5
                            direction = 1
                    else:
                        feedback = "Fix Form"

                if per == 100:
                    if elbow > 160 and shoulder > 40 and hip > 160:
                        feedback = "Down"
                        if direction == 1:
                            count += 0.5
                            direction = 0
                    else:
                        feedback = "Fix Form"
                        # form = 0

            # Draw Bar and Progress Bar
            if form == 1:
                cv2.rectangle(img, (580, 50), (600, 380), (255, 0, 0), 3)
                cv2.rectangle(img, (580, int(bar)), (600, 380), (255, 0, 0), cv2.FILLED)
                cv2.putText(img, f'{int(per)}%', (565, 430), cv2.FONT_HERSHEY_PLAIN, 1,
                            (0, 0, 0), 2)

                # Add Progress Bar for Center of Gravity
                cog = calculate_center_of_gravity(lmList)
                x, y = cog
                balance = calculate_body_balance(lmList[11:25])  # 取得身体关键点的数据
                distance_to_balance = np.linalg.norm(cog - np.array(balance))  # 计算距离
                min_distance = 0  # 最小距离
                max_distance = 50  # 最大距离
                min_x = -5.0  # 缩放范围最小值
                max_x = 5.0  # 缩放范围最大值

                progress_position = int(np.interp(distance_to_balance, (min_distance, max_distance), (min_x, max_x))+5)

                if distance_to_balance <= 50:  # 判断是否在推荐的重心范围内
                    balance_feedback = "In Balance"
                else:
                    balance_feedback = "Out of Balance"
                progress_bar = ""
                for i in range(10):
                    if i == progress_position:
                        progress_bar += "|"
                    else:
                        progress_bar += "-"
                cv2.putText(img, 'CoG:', (20, 40), cv2.FONT_HERSHEY_PLAIN, 1,
                            (0, 0, 0), 2)
                cv2.putText(img, progress_bar, (100, 40), cv2.FONT_HERSHEY_PLAIN, 1,
                            (0, 0, 0), 2)
                cv2.putText(img, '({:.2f}, {:.2f})'.format(x, y), (220, 40), cv2.FONT_HERSHEY_PLAIN, 1,
                            (0, 0, 0), 2)
                cv2.putText(img, balance_feedback, (350, 40), cv2.FONT_HERSHEY_PLAIN, 1,
                            (0, 0, 255), 2)

            # Pushup counter
            cv2.rectangle(img, (0, 380), (100, 480), (255, 255, 255), cv2.FILLED)
            cv2.putText(img, str(int(count)), (25, 455), cv2.FONT_HERSHEY_SIMPLEX, 2,
                        (0, 0, 0), 5)

            # Feedback
            cv2.rectangle(img, (500, 0), (640, 40), (255, 255, 255), cv2.FILLED)
            cv2.putText(img, feedback, (500, 30), cv2.FONT_HERSHEY_SIMPLEX, 1,
                        (0, 0, 0), 2)

        info["pushup_count"] = count
        _, buffer = cv2.imencode('.jpg', img)
        frame = buffer.tobytes()
        yield frame, info
    cap.release()


@app.route('/video_feed')
def video_feed():
    def inner():
        for frame, _ in generate_frames():
            yield (b'--frame\r\n'
                   b'Content-Type: image/jpeg\r\n\r\n' + frame + b'\r\n')
    return Response(inner(), mimetype='multipart/x-mixed-replace; boundary=frame')


@app.route('/data')
def data():
    _, info = next(generate_frames())
    return jsonify(info)


if __name__ == "__main__":
    app.run(debug=True)
